feat(zmq): multimodal inputs over the direct ZMQ backend (3/7) - #2056

crates/engine_zmq_client/src/protocol/vllm/multimodal.rs (1)

248-260: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Add MmField::Flat to the round-trip test.

field_round_trips_python_factory_tuple covers only Batched and Shared. Flat is the variant whose MmFieldWireInner untagged resolution depends on deny_unknown_fields distinguishing three field sets, and it carries the nested MmSlice/SliceSpec tuple encoding. That is the most fragile decode path in this module and it is only exercised on the encode side.

♻️ Add the Flat case

         for field in [
             MmField::Batched(MmBatchedField { keep_on_cpu: true }),
             MmField::Shared(MmSharedField {
                 batch_size: 4,
                 keep_on_cpu: false,
             }),
+            MmField::Flat(MmFlatField {
+                slices: vec![MmSlice::Slice(SliceSpec {
+                    start: Some(0),
+                    stop: Some(1200),
+                    step: None,
+                })],
+                dim: 0,
+                keep_on_cpu: true,
+            }),
         ] {

As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs` around lines 248 -
260, Add the MmField::Flat variant, including representative MmSlice/SliceSpec
data, to the field_round_trips_python_factory_tuple test so its nested tuple
encoding and untagged decoding are verified alongside Batched and Shared.

Source: Coding guidelines

model_gateway/src/routers/grpc/backend_client.rs (1)

368-379: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

🟡 Nit: zmq_tokenspeed_mm is unreachable today, and the into_proto(true) flag is unexplained.

Two other layers already reject TokenSpeed multimodal input: assemble_multimodal_data_impl bails for a non-vLLM ZMQ runtime, and translate_request_tokenspeed rejects any request with mm_inputs set. So this converter can only ever produce mm_inputs that the translator immediately rejects with a generic message. Add a short comment stating that the converter exists for the future TokenSpeed wire, so a reader does not assume the path works.

Also document what true selects in into_proto(true), or use a named binding. The literal carries no meaning at the call site.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/routers/grpc/backend_client.rs` around lines 368 - 379, Add
a brief comment above zmq_tokenspeed_mm explaining that it is retained for a
future TokenSpeed wire path despite being currently unreachable. Replace the
unexplained true argument in MultimodalData::TokenSpeed conversion with a named
binding or concise explanation identifying the serialization mode it selects,
without changing behavior.

model_gateway/src/routers/grpc/zmq_multimodal.rs (1)

384-405: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Two coverage gaps, and a fixture that hides the bounds check.

• base_inputs() uses placeholders (1, 3) and (6, 3), but several tests pass &[] as prompt_token_ids. Those ranges exceed an empty prompt. The tests pass only because is_embed_mask returns early when im_token_id is None. If you apply the bounds-check fix, give these tests a prompt of at least 9 tokens.
• No test sets keep_on_cpu_keys, so the dropped on_cpu in the shared branch is not caught.

As per coding guidelines: "Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality."

Also applies to: 447-469

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/routers/grpc/zmq_multimodal.rs` around lines 384 - 405,
Update tests using base_inputs and empty prompt_token_ids to provide a prompt
with at least nine tokens, so placeholder ranges are valid once bounds checks
are enforced. Add coverage for keep_on_cpu_keys through the shared
feature-building path and assert the resulting on-CPU behavior is preserved. Run
the pr-test-analyzer agent to verify the new and changed functionality is
adequately covered.

Source: Coding guidelines

model_gateway/src/routers/grpc/multimodal/assemble.rs (1)

114-127: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

🟡 Nit: The ZMQ runtime dispatch shape differs from backend_client.rs.

This arm treats Vllm | Unspecified as vLLM and bails for anything else. BackendClient::build_chat_request and build_messages_request in model_gateway/src/routers/grpc/backend_client.rs (lines 209-236, 254-281) use the opposite shape: TokenSpeed is explicit and _ falls through to vLLM. Both functions switch on the same client.runtime(). Today ZmqEngineClient::runtime() returns only Vllm or TokenSpeed, so no live path diverges. If a third ZMQ runtime is added, assembly will reject the request while request building silently treats it as vLLM. Pick one shape for both.

Separately, assemble_vllm computes shm_enabled and shm_min_bytes, and this arm then overwrites shm_enabled and rdma_enabled. shm_min_bytes is left at its resolved value even though it no longer applies. Passing the inline-transport intent into assemble_vllm would keep the three transport fields consistent in one place.

As per coding guidelines: "Configuration changes must be checked across CLI arguments, types.rs, both conversion paths in main.rs, Python bindings, and the Go SDK."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model_gateway/src/routers/grpc/multimodal/assemble.rs` around lines 114 -
127, Align the ZMQ runtime dispatch in the BackendClient::Zmq arm with
BackendClient::build_chat_request and build_messages_request, keeping TokenSpeed
explicit and routing other runtimes to vLLM. Update assemble_vllm to receive the
inline-transport intent so it sets shm_enabled, rdma_enabled, and shm_min_bytes
consistently, removing the post-assembly field overrides.

Source: Coding guidelines

crates/engine_zmq_client/src/codec/tensor.rs (1)

158-159: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Reuse the existing float32 payload instead of copying it.

from_f32_bytes_cast already owns the allocation after data.len().is_multiple_of(4) and as_chunks::<4> in crates/engine_zmq_client/src/codec/tensor.rs; the Float32 branch still calls data.to_vec(), so decode_tensor immediately discards the original Vec<u8> and allocates a second copy. Allow from_raw_bytes/from_raw to accept or take Bytes so the no-cast path can reuse the inline payload for large pixel tensors.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/engine_zmq_client/src/codec/tensor.rs` around lines 158 - 159, Update
the Float32 branch in decode_tensor and the related from_raw_bytes/from_raw
constructors to accept or take ownership of Bytes, reusing the existing
from_f32_bytes_cast allocation instead of calling data.to_vec(). Preserve the
current validation and decoding behavior while eliminating the redundant payload
copy.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

crates/engine_zmq_client/src/protocol/vllm/multimodal.rs (1)

247-259: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

🟡 Nit: Add round-trip coverage for MmField::Flat.

The test only round-trips Batched and Shared. Flat uses the distinct slices tuple encoding and factory conversion path. Add Flat cases for both MmSlice::Slice and MmSlice::Slices.

As per coding guidelines, “Run the pr-test-analyzer agent to verify that tests adequately cover new or changed functionality.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/engine_zmq_client/src/protocol/vllm/multimodal.rs` around lines 247 -
259, Extend field_round_trips_python_factory_tuple to include MmField::Flat
cases containing both MmSlice::Slice and MmSlice::Slices, then verify each
encodes and decodes back identically through the existing round-trip assertions.
Run the pr-test-analyzer agent to confirm coverage.

Source: Coding guidelines